from src.tools.dms_operator import dmsoperator
from src.tools.mml_to_latex import mathml2latex
from src.edit_file.to_document import to_document_sop
from src.edit_file.db import faiss_embedding1, faiss_embedding2
from src.edit_file.prompt import Prompt
from src.tools.chain import gen_chain
from src.tools.split_string import split_sub_string
from src.tools.extract_jsonstring import re_json_string
from src.tools.information_exchange import pub_gen
import json

class EditSOPFile:

    def __init__(self):
        self.prompt = Prompt()
        self.insertsort = dict()

    # 解析表格
    def parse_table(self, table: dict):
        table_dict = {
            "index": table["index"],
            "sectionIndex": table["sectionIndex"],
            "bodyChildIndex": table["bodyChildIndex"],
            "tableCellInfos": []
        }
        table_text = ""
        for cell in table["tableCellInfos"]:
            cell_dict = {
                "rowIndex": cell["rowIndex"],
                "columnIndex": cell["columnIndex"],
                "paragraphInfos": []
            }
            cell_text = ""
            for para in cell["paragraphInfos"]:
                para_text = ""
                for textinfo in para["textInfos"]:
                    if textinfo["subSuperscript"] == 2:
                        para_text += f'<sub>{textinfo["text"]}</sub>'
                    else:
                        para_text += textinfo["text"]

                mathinfos = []
                for mathinfo in para["officeMathInfos"]:
                    math_dict = {
                        "text": mathml2latex(mathinfo["mathMlCode"])
                    }
                    mathinfos.append(math_dict)

                para_dict = {
                    "index": para["index"],
                    "sectionIndex": para["sectionIndex"],
                    "bodyChildIndex": para["bodyChildIndex"],
                    "text": para_text,
                    "latexMathInfos": mathinfos
                }
                cell_dict["paragraphInfos"].append(para_dict)
                cell_text += para["text"]
            table_dict["tableCellInfos"].append(cell_dict)
            table_text += cell_text + "|"
        return table_dict, table_text

    def edit_table(self, revise_SOPtext_dict, table, file_pubinfo):
        if len(revise_SOPtext_dict["tableCellInfos"]) == len(table["tableCellInfos"]):
            table_dict = {
                "index": revise_SOPtext_dict["index"],
                "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                "operationType": 0,
                "tableInfo": {
                    "cellEditInfos": []
                }
            }
            for cell in revise_SOPtext_dict["tableCellInfos"]:
                cell_dict = {
                    "rowIndex": cell["rowIndex"],
                    "columnIndex": cell["columnIndex"],
                    "cellChildInfos": []
                }
                for para in cell["paragraphInfos"]:
                    edit_para_dict = {
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "operationType": 1,
                        "paragraphInfo": {
                            "paragraphChildInfos": [],
                            "referenceParagraphType": 2
                        }
                    }
                    # 插入纯文本
                    textinfo_list = split_sub_string(para["text"])
                    textinfos = []
                    for textinfo_index, textinfo in enumerate(textinfo_list):
                        textinfo_dict = {
                            "paragraphChildIndex": textinfo_index,
                            "paraChildType": 0,
                            "content": textinfo.get("content"),
                            "subSuperScript": textinfo.get("subSuperScript")
                        }
                        textinfos.append(textinfo_dict)
                    if textinfos:
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                    # 插入数学公式
                    if para["latexMathInfos"]:
                        latex_index = len(textinfos)
                        mathinfos = []
                        for latex_math_info in para["latexMathInfos"]:
                            mathinfo_dict = {
                                "paragraphChildIndex": latex_index,
                                "paraChildType": 2,
                                "content": latex_math_info["text"]
                            }
                            latex_index += 1
                            mathinfos.append(mathinfo_dict)
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
                    cell_dict["cellChildInfos"].append(edit_para_dict)
                    # 删除段落
                    dele_para_dict = {
                        "index": para["index"],
                        "sectionIndex": para["sectionIndex"],
                        "bodyChildIndex": para["bodyChildIndex"],
                        "operationType": 2
                    }
                    cell_dict["cellChildInfos"].append(dele_para_dict)
                table_dict["tableInfo"]["cellEditInfos"].append(cell_dict)
            file_pubinfo["editInfos"].append(table_dict)

        else:
            # 先插入表格
            max_row = -1
            max_column = -1
            for cell in revise_SOPtext_dict["tableCellInfos"]:
                max_row = max(cell["rowIndex"], max_row)
                max_column = max(cell["columnIndex"], max_column)
            row_num = max_row + 1
            column_num = max_column + 1
            table_dict = {
                "index": revise_SOPtext_dict["index"],
                "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                "operationType": 1,
                "tableInfo": {
                    "tableInsertInfo": {
                        "rowCount": row_num,
                        "columnCount": column_num,
                        "cellInsertInfos": []
                    }
                }
            }
            for cell in revise_SOPtext_dict["tableCellInfos"]:
                cell_dict = {
                    "rowIndex": cell["rowIndex"],
                    "columnIndex": cell["columnIndex"],
                    "cellInsertContentInfos": []
                }
                for para in cell["paragraphInfos"]:

                    edit_para_dict = {
                        "index": para["index"],
                        "paragraphInfo": {
                            "paragraphChildInfos": [],
                        }
                    }
                    # 先插入纯文本
                    textinfo_list = split_sub_string(para["text"])
                    textinfos = []
                    for textinfo_index, textinfo in enumerate(textinfo_list):
                        textinfo_dict = {
                            "paragraphChildIndex": textinfo_index,
                            "paraChildType": 0,
                            "content": textinfo.get("content"),
                            "subSuperScript": textinfo.get("subSuperScript")
                        }
                        textinfos.append(textinfo_dict)
                    if textinfos:
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                    # 插入数学公式
                    if para["latexMathInfos"]:
                        latex_index = len(textinfos)
                        mathinfos = []
                        for latex_math_info in para["latexMathInfos"]:
                            mathinfo_dict = {
                                "paragraphChildIndex": latex_index,
                                "paraChildType": 2,
                                "content": latex_math_info["text"]
                            }
                            latex_index += 1
                            mathinfos.append(mathinfo_dict)
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
                    cell_dict["cellInsertContentInfos"].append(edit_para_dict)
                table_dict["tableInfo"]["tableInsertInfo"]["cellInsertInfos"].append(cell_dict)
            file_pubinfo["editInfos"].append(table_dict)
            # 再删除表格
            dele_table = {
                "index": revise_SOPtext_dict["index"],
                "sectionIndex": revise_SOPtext_dict["sectionIndex"],
                "bodyChildIndex": revise_SOPtext_dict["bodyChildIndex"],
                "operationType": 2,
            }
            file_pubinfo["editInfos"].append(dele_table)

    def insert_table(self, revise_SOPtext_dict, table, file_pubinfo, sortnum):
        max_row = -1
        max_column = -1
        for cell in revise_SOPtext_dict["tableCellInfos"]:
            max_row = max(cell["rowIndex"], max_row)
            max_column = max(cell["columnIndex"], max_column)
        row_num = max_row + 1
        column_num = max_column + 1
        table_dict = {
            "index": table["index"],
            "sectionIndex": table["sectionIndex"],
            "bodyChildIndex": table["bodyChildIndex"] + 1,
            "sort": sortnum,
            "operationType": 1,
            "tableInfo": {
                "tableInsertInfo": {
                    "rowCount": row_num,
                    "columnCount": column_num,
                    "cellInsertInfos": []
                }
            }
        }
        for cell in revise_SOPtext_dict["tableCellInfos"]:
            cell_dict = {
                "rowIndex": cell["rowIndex"],
                "columnIndex": cell["columnIndex"],
                "cellInsertContentInfos": []
            }
            for para in cell["paragraphInfos"]:

                edit_para_dict = {
                    "index": para["index"],
                    "paragraphInfo": {
                        "paragraphChildInfos": [],
                    }
                }
                # 先插入纯文本
                textinfo_list = split_sub_string(para["text"])
                textinfos = []
                for textinfo_index, textinfo in enumerate(textinfo_list):
                    textinfo_dict = {
                        "paragraphChildIndex": textinfo_index,
                        "paraChildType": 0,
                        "content": textinfo.get("content"),
                        "subSuperScript": textinfo.get("subSuperScript")
                    }
                    textinfos.append(textinfo_dict)
                if textinfos:
                    edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                # 插入数学公式
                if para["latexMathInfos"]:
                    latex_index = len(textinfos)
                    mathinfos = []
                    for latex_math_info in para["latexMathInfos"]:
                        mathinfo_dict = {
                            "paragraphChildIndex": latex_index,
                            "paraChildType": 2,
                            "content": latex_math_info["text"]
                        }
                        latex_index += 1
                        mathinfos.append(mathinfo_dict)
                    edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += mathinfos
                cell_dict["cellInsertContentInfos"].append(edit_para_dict)
            table_dict["tableInfo"]["tableInsertInfo"]["cellInsertInfos"].append(cell_dict)
        file_pubinfo["editInfos"].append(table_dict)

    # 找到流式输出内容
    def find_key_values(self, target_key, data):
        found_values = []
        if isinstance(data, dict):
            for key, value in data.items():
                if key == target_key and value:
                    found_values.append(value)
                found_values.extend(self.find_key_values(target_key, value))
        elif isinstance(data, list):
            for item in data:
                found_values.extend(self.find_key_values(target_key, item))
        return found_values

    def edit_file(self, revise_file, template_file_path, task_id):
        revise_file_path = revise_file.get("realPath")
        revise_jsonfile = dmsoperator.get_jsonfile(revise_file_path)
        # print("-" * 200)
        # file_pubinfo_json = json.dumps(revise_jsonfile, indent=4, ensure_ascii=False)
        # print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"解析.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-" * 200)
        # 要修订文档入向量库
        docus_para, docus_table = to_document_sop(revise_jsonfile, revise_file)
        file_pubinfo = {
            "filePath": revise_jsonfile["filePath"],
            "isHighlight": False,
            "editInfos": []
        }
        # 修改段落
        faiss_embedding1.add_documents(docus_para)
        retr_para = faiss_embedding1.as_retriever(search_kwargs={"k": 1})
        template_diffjson = dmsoperator.get_diff_jsonfile(template_file_path)
        # print("-" * 200)
        # file_pubinfo_json = json.dumps(template_diffjson, indent=4, ensure_ascii=False)
        # print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"段落修改.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-" * 200)
        for diff_chunk in template_diffjson["data"]:
            for diff_content in diff_chunk["detailDiffDtoList"]:
                # 编辑类型
                if diff_content["operationType"] == 0:
                    revise_retrpara_list = retr_para.invoke(diff_content["beforeContent"])
                    revise_retrpara = revise_retrpara_list[0]
                    revise_query = self.prompt.edit_para_prompt(diff_content["beforeContent"], diff_content["afterContent"], revise_retrpara.page_content)
                    revise_ans = gen_chain(revise_query)
                    pub_gen(task_id, revise_ans, revise_file.get("docId"), revise_file.get("taskDocId"))
                    edit_para_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"],
                        "operationType": 1,
                        "paragraphInfo": {
                            "paragraphChildInfos": [],
                            "referenceParagraphType": 2
                        }
                    }
                    # 插入纯文本
                    textinfo_list = split_sub_string(revise_ans)
                    textinfos = []
                    for textinfo_index, textinfo in enumerate(textinfo_list):
                        textinfo_dict = {
                            "paragraphChildIndex": textinfo_index,
                            "paraChildType": 0,
                            "content": textinfo.get("content"),
                            "subSuperScript": textinfo.get("subSuperScript")
                        }
                        textinfos.append(textinfo_dict)
                    if textinfos:
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                    file_pubinfo["editInfos"].append(edit_para_dict)
                    # 删除段落
                    dele_edit_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"],
                        "operationType": 2
                    }
                    file_pubinfo["editInfos"].append(dele_edit_dict)
                # 插入类型
                if diff_content["operationType"] == 1:
                    revise_retrpara_list = retr_para.invoke(diff_chunk["previousParagraph"])
                    revise_retrpara = revise_retrpara_list[0]
                    # 判断插入顺序
                    if diff_chunk["previousParagraph"] not in self.insertsort:
                        self.insertsort[diff_chunk["previousParagraph"]] = 0
                    else:
                        self.insertsort[diff_chunk["previousParagraph"]] += 1
                    edit_para_dict = {
                        "index": revise_retrpara.metadata["index"],
                        "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                        "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"] + 1,
                        "sort": self.insertsort[diff_chunk["previousParagraph"]],
                        "operationType": 1,
                        "paragraphInfo": {
                            "paragraphChildInfos": [],
                        }
                    }
                    # 插入图片
                    if "<KLPicture>" in diff_content["afterContent"]:
                        picinfo_dict = {
                            "paragraphChildIndex": 0,
                            "paraChildType": 1,
                            "content": diff_content["afterContent"].removeprefix("<KLPicture>").removesuffix("</KLPicture>"),
                            "subSuperScript": 0
                        }
                        edit_para_dict["paragraphInfo"]["paragraphChildInfos"].append(picinfo_dict)
                    # 插入纯文本
                    else:
                        pub_gen(task_id, diff_content["afterContent"], revise_file.get("docId"), revise_file.get("taskDocId"))
                        textinfo_list = split_sub_string(diff_content["afterContent"])
                        textinfos = []
                        for textinfo_index, textinfo in enumerate(textinfo_list):
                            textinfo_dict = {
                                "paragraphChildIndex": textinfo_index,
                                "paraChildType": 0,
                                "content": textinfo.get("content"),
                                "subSuperScript": textinfo.get("subSuperScript")
                            }
                            textinfos.append(textinfo_dict)
                        if textinfos:
                            edit_para_dict["paragraphInfo"]["paragraphChildInfos"] += textinfos
                    file_pubinfo["editInfos"].append(edit_para_dict)

                # 删除类型
                if diff_content["operationType"] == 2:
                    # 删除图片
                    if "<KLPicture>" in diff_content["beforeContent"]:
                        revise_retrpara_list = retr_para.invoke(diff_chunk["previousParagraph"])
                        print(revise_retrpara_list)
                        revise_retrpara = revise_retrpara_list[0]
                        dele_edit_dict = {
                            "index": revise_retrpara.metadata["index"] + 1,
                            "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                            "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"] + 1,
                            "operationType": 2
                        }
                        file_pubinfo["editInfos"].append(dele_edit_dict)
                    # 删除文字
                    else:
                        revise_retrpara_list = retr_para.invoke(diff_content["beforeContent"])
                        revise_retrpara = revise_retrpara_list[0]
                        dele_edit_dict = {
                            "index": revise_retrpara.metadata["index"],
                            "sectionIndex": revise_retrpara.metadata["sectionIndex"],
                            "bodyChildIndex": revise_retrpara.metadata["bodyChildIndex"],
                            "operationType": 2
                        }
                        file_pubinfo["editInfos"].append(dele_edit_dict)

        # 修改表格
        faiss_embedding2.add_documents(docus_table)
        retr_table = faiss_embedding2.as_retriever(search_kwargs={"k": 1})
        template_diffjson = dmsoperator.get_diff_jsonfile_table(template_file_path)
        # print("-"*200)
        # file_pubinfo_json = json.dumps(template_diffjson, indent=4, ensure_ascii=False)
        # print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"表格修改.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-"*200)
        for diff_chunk in template_diffjson["data"]:
            if diff_chunk["tableDiffInfo"] is None:
                continue
            table_diffinfo = diff_chunk["tableDiffInfo"]
            # 编辑表格
            if table_diffinfo["operationType"] == 0:
                beforetable_dict, beforetable_text = self.parse_table(table_diffinfo["beforeTableInfo"])
                aftertable_dict, aftertable_text = self.parse_table(table_diffinfo["afterTableInfo"])
                revise_retrtable_list = retr_table.invoke(beforetable_text)
                revise_retrtable = revise_retrtable_list[0]
                beforetable_jsontext = json.dumps(beforetable_dict, ensure_ascii=False)
                aftertable_jsontext = json.dumps(aftertable_dict, ensure_ascii=False)
                revise_retrtable_jsontext = json.dumps(revise_retrtable.metadata, ensure_ascii=False)
                revise_query = self.prompt.edit_table_prompt(beforetable_jsontext, aftertable_jsontext, revise_retrtable_jsontext)
                revise_ans = gen_chain(revise_query)
                revise_ans_textjson = re_json_string(revise_ans)
                revise_ans_dict = json.loads(revise_ans_textjson)
                print(f"编辑表格{revise_ans_dict}")
                text_key = self.find_key_values("text", revise_ans_dict)
                for msg in text_key:
                    pub_gen(task_id, msg, revise_file.get("docId"), revise_file.get("taskDocId"))
                self.edit_table(revise_ans_dict, revise_retrtable.metadata, file_pubinfo)

            # 插入表格
            if table_diffinfo["operationType"] == 1:
                if table_diffinfo["previousParagraph"] is None:
                    continue
                revise_retrpara_list = retr_para.invoke(table_diffinfo["previousParagraph"])
                revise_retrpara = revise_retrpara_list[0]
                aftertable_dict, aftertable_text = self.parse_table(table_diffinfo["afterTableInfo"])
                text_key = self.find_key_values("text", aftertable_dict)
                for msg in text_key:
                    pub_gen(task_id, msg, revise_file.get("docId"), revise_file.get("taskDocId"))
                    # 判断插入顺序
                if table_diffinfo["previousParagraph"] not in self.insertsort:
                    self.insertsort[table_diffinfo["previousParagraph"]] = 0
                else:
                    self.insertsort[table_diffinfo["previousParagraph"]] += 1
                self.insert_table(aftertable_dict, revise_retrpara.metadata, file_pubinfo, self.insertsort[table_diffinfo["previousParagraph"]])
            # 删除表格
            if table_diffinfo["operationType"] == 2:
                beforetable_dict, beforetable_text = self.parse_table(table_diffinfo["beforeTableInfo"])
                revise_retrtable_list = retr_table.invoke(beforetable_text)
                revise_retrtable = revise_retrtable_list[0]
                dele_table = {
                    "index": revise_retrtable.metadata["index"],
                    "sectionIndex": revise_retrtable.metadata["sectionIndex"],
                    "bodyChildIndex": revise_retrtable.metadata["bodyChildIndex"],
                    "operationType": 2,
                }
                file_pubinfo["editInfos"].append(dele_table)
        # print("-" * 200)
        # file_pubinfo_json = json.dumps(file_pubinfo, indent=4, ensure_ascii=False)
        # print(file_pubinfo_json)
        # # 保存修改的json文件
        # edit_SOPjsonfile_path = f"修改.json"
        # with open(edit_SOPjsonfile_path, "w", encoding="utf-8") as f:
        #     f.write(file_pubinfo_json)
        # print("-"*200)

        edit_res = dmsoperator.post_edit_file_v2(file_pubinfo)
        res = dmsoperator.post_compair(revise_file.get("realPath"), edit_res["data"])
        return res["data"]





